package stablediffusion

import (
	httpUtil "app/pkg/http"
	"app/pkg/log"
	"app/pkg/oss"
	"encoding/json"
	"errors"
	"fmt"
	"time"
)

type LibClient struct {
	BaseUrl   string
	Token     string
	Referer   string
	UserAgent string
}

// 选择lora模型时的配置
type AdditionalNetwork struct {
	ModelId          int64   `json:"modelId"`
	Type             int8    `json:"type"`
	Weight           float32 `json:"weight"`
	ModelName        string  `json:"modelName"`
	ModelVersionName string  `json:"modelVersionName"`
}

type LibText2imgParams struct {
	Prompt         string `json:"prompt"`
	NegativePrompt string `json:"negativePrompt"`
	ExtraNetwork   string `json:"extraNetwork"`
	SamplingMethod int    `json:"samplingMethod"`
	SamplingStep   int    `json:"samplingStep"`
	Width          int    `json:"width"`
	Height         int    `json:"height"`
	BatchCount     int    `json:"batchCount"`
	BatchSize      int    `json:"batchSize"`
	CfgScale       int    `json:"cfgScale"`
	Seed           int    `json:"seed"`
	SeedExtra      int    `json:"seedExtra"`
	HiResFix       int    `json:"hiResFix"`
	RestoreFaces   int    `json:"restoreFaces"`
	Tiling         int    `json:"tiling"`
}

type LibGenerateImageParams struct {
	CheckpointId      int64                  `json:"checkpointId"`
	GenerateType      int                    `json:"generateType"` //1
	FrontCustomerReq  map[string]interface{} `json:"frontCustomerReq"`
	Text2img          LibText2imgParams      `json:"text2img"`
	AdetailerEnable   int                    `json:"adetailerEnable"`
	AdditionalNetwork []AdditionalNetwork    `json:"additionalNetwork"`
}

type LibGenerateImageReply struct {
	Code int    `json:"code"`
	Data int64  `json:"data"`
	Msg  string `json:"msg"`
}

func NewLibClient() *LibClient {
	return &LibClient{
		BaseUrl:   "https://www.liblibai.com",
		Referer:   "https://www.liblibai.com/editor/",
		UserAgent: "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/115.0.0.0 Safari/537.36",
	}
}

func (m *LibClient) WithToken(t string) {
	m.Token = t
}

func (m *LibClient) WithReferer(r string) {
	m.Referer = r
}

func (m *LibClient) WithUserAgent(r string) {
	m.UserAgent = r
}
func (m *LibClient) getHeaders() (h map[string]string, err error) {
	h = make(map[string]string)
	if m.UserAgent == "" {
		err = fmt.Errorf("lib user agent not set")
		return
	}
	h["User-Agent"] = m.UserAgent

	if m.Token == "" {
		err = fmt.Errorf("lib token not set")
		return
	}
	h["Token"] = m.Token

	if m.Referer == "" {
		err = fmt.Errorf("lib referer not set")
		return
	}
	h["Referer"] = m.Referer

	return
}

// sd 接口参数转换成 lib txt2img接口参数
func (m *LibClient) sdTxt2imgParamsTrans(txt2Img Txt2ImgRequest) (params LibText2imgParams) {
	params.Prompt = txt2Img.Prompt
	params.NegativePrompt = txt2Img.NegativePrompt
	params.SamplingStep = txt2Img.Steps
	params.Width = txt2Img.Width
	params.Height = txt2Img.Height
	params.Seed = txt2Img.Seed
	params.BatchCount = txt2Img.BatchSize
	params.BatchSize = txt2Img.BatchSize
	params.CfgScale = txt2Img.CfgScale
	return
}

type LibGenerateProgressReply struct {
	Code int `json:"code"`
	Data struct {
		PercentCompleted int `json:"percentCompleted"`
		Images           []struct {
			PreviewPath string `json:"previewPath"`
		} `json:"images"`
	} `json:"data"`
	Msg string `json:"msg"`
}

// libId 对应lib checkpointId
func (m *LibClient) Txt2ImgWithProgress(libId int64, txt2Img Txt2ImgRequest, network []AdditionalNetwork, ch chan int64) (pics []string, downErrPics []string, libTaskId int64, err error) {
	headers, err := m.getHeaders()
	if err != nil {
		return
	}
	//构造参数
	req := LibGenerateImageParams{
		CheckpointId: libId,
		GenerateType: 1,
		FrontCustomerReq: map[string]interface{}{
			"frontId":         "79995446-3a98-42ff-a35b-30fb6991b646",
			"windowId":        "",
			"tabType":         "txt2img",
			"conAndSegAndGen": "gen",
		},
		Text2img:          m.sdTxt2imgParamsTrans(txt2Img),
		AdditionalNetwork: network,
	}
	isEnd := false

	log.NewDebug("sd-api/generate/image req: ", fmt.Sprintf("%+v", req))
	res, err := httpUtil.JsonPost(fmt.Sprintf("%s/gateway/sd-api/generate/image", m.BaseUrl), req, headers, 20)
	if err != nil {
		return
	} else {
		reply := LibGenerateImageReply{}
		err = json.Unmarshal(res, &reply)
		if err != nil {
			return
		}
		if reply.Data <= 0 {
			err = errors.New(reply.Msg)
			return
		}
		libTaskId = reply.Data
		count := 0
		// 这里更新进度
		for !isEnd {
			res, err := m.Progress(libTaskId)
			if err != nil || count > 300 { //300秒后放弃
				log.NewError("lib txt2image Progress Error: ", err.Error())
				isEnd = true
			}
			log.NewInfo("lib progress: ", fmt.Sprintf("%+v", res))
			if res.Data.PercentCompleted != 100 {
				ch <- int64(res.Data.PercentCompleted)
				count += 1
			} else {
				isEnd = true
				for _, v := range res.Data.Images {
					//将图片上传到七牛云
					qnOss := oss.NewQnOss()
					headers := map[string]string{
						"Referer":            "https://www.liblibai.com/",
						"Sec-Fetch-Dest":     "image",
						"Connection":         "keep-alive",
						"Origin":             "https://www.liblibai.com",
						"Sec-Fetch-Mode":     "cors",
						"Sec-Fetch-Site":     "same-site",
						"sec-ch-ua-mobile":   "?0",
						"sec-ch-ua-platform": "macOS",
						"sec-ch-ua":          `Not/A)Brand";v="99", "Google Chrome";v="115", "Chromium";v="115`,
						"User-Agent":         "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/115.0.0.0 Safari/537.36",
						"Accept-Encoding":    "gzip, deflate, br",
						"Accept-Language":    "zh-CN,zh;q=0.9,en;q=0.8",
						"Host":               "liblibai-tmp-image.liblibai.com",
						"Accept":             "image/avif,image/webp,image/apng,image/svg+xml,image/*,*/*;q=0.8",
					}
					pic, err := qnOss.PutNetImage(v.PreviewPath, headers)
					if err != nil {
						log.NewError("PutNetImage Error", err.Error(), fmt.Sprintf("%+v", v))
						if err.Error() == "error http code 403" {
							//重试一次 不行就发邮件手动上传
							pic, err = qnOss.PutNetImage(v.PreviewPath, headers)
							if err != nil {
								log.NewError("PutNetImage Error", err.Error(), fmt.Sprintf("%+v", v))
								downErrPics = append(downErrPics, v.PreviewPath)
								continue
							}
						} else {
							downErrPics = append(downErrPics, v.PreviewPath)
							continue
						}
					}
					pics = append(pics, pic)
				}
			}
		}
		close(ch)

	}
	return
}

func (m *LibClient) Progress(taskId int64) (reply LibGenerateProgressReply, err error) {
	headers, err := m.getHeaders()
	if err != nil {
		return
	}
	res, err := httpUtil.JsonPost(fmt.Sprintf("%s/api-websocket/gen/img/progress/msg/%d", m.BaseUrl, taskId), map[string]interface{}{}, headers, 60)
	if err != nil {
		return
	}
	err = json.Unmarshal(res, &reply)
	time.Sleep(1 * time.Second)
	return
}

type TokenUsageReply struct {
	Code int `json:"code"`
	Data struct {
		DailyThreshold int `json:"dailyThreshold"`
		UsedCount      int `json:"usedCount"`
	} `json:"data"`
	Msg string `json:"msg"`
}

func (m *LibClient) TokenUsage() (reply TokenUsageReply, err error) {
	headers, err := m.getHeaders()
	if err != nil {
		return
	}
	res, err := httpUtil.Get(fmt.Sprintf("%s/gateway/sd-api/generate/usage", m.BaseUrl), headers, 60)
	if err != nil {
		return
	}
	err = json.Unmarshal(res, &reply)
	return
}
